package com.example.spring_ai.a2a.server.controller;

import com.example.spring_ai.a2a.server.bean.TaskRequest;
import com.example.spring_ai.a2a.server.bean.TaskResponse;
import com.example.spring_ai.a2a.server.service.TaskService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;

/**
 * TaskController
 * 实现任务提交与异步处理逻辑
 *
 **/
@RestController
@Slf4j
public class AgentController {

    @Autowired
    private TaskService taskService;

    @GetMapping("/.well-known/agent.json")
    public Map<String, Object> getAgentCard() {
        return Map.of(
                "name", "CalculateAgent",
                "version", "1.0",
                "description", "提供计算服务",
                "endpoints", Map.of(
                        "task_send", "/api/tasks/submit",
                        "task_get", "/api/tasks/subscribe_sse"
                ),
                "input_schema", Map.of(
                        "type", "object",
                        "properties", Map.of(
                                "tasks", Map.of("type", "string", "enum", List.of("task1", "task2", "task3")),
                                "required", List.of("tasks")
                        ),
                        "authentication", Map.of("methods", List.of("API_Key"))
                ));
    }

    @PostMapping("/api/tasks/submit")
    public ResponseEntity<TaskResponse> submitTask(
            @RequestHeader("Authorization") String apiKey,
            @RequestBody TaskRequest request
    ) {
        // 认证校验（示例简化）
        if (!"SECRET_KEY".equals(apiKey)) {
            return ResponseEntity.status(401).build();
        }

        // 异步处理任务
        CompletableFuture.runAsync(() ->
                taskService.processTask(request.getTaskId(), request.getParams())
        );

        return ResponseEntity.accepted().body(
                new TaskResponse(request.getTaskId(), "submitted", "", "")
        );
    }

    @GetMapping(value = "/api/tasks/subscribe_sse", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<TaskResponse> streamTaskStatus(@RequestParam("taskId") String taskId) {
        log.info("[server] updates task {} with streamTaskStatus", taskId);
        return taskService.streamTaskStatus(taskId);
    }

}
